/*--------------------------------------------------------------------------------------------------
Yatak Operation System Project Copyright 2012-2013
Author : Ilya Churaev
Started : 17 Jun 2012		Updated : 12 Feb 2013
File Name : kheap.c		Purpose : Interface for kernel heap functionsm.
--------------------------------------------------------------------------------------------------*/
#include <yatak/paging.h>
#include <yatak/mutex.h>
#include <yatak/kheap.h>
#include <yatak/yatak.h>
#include <screen.h>

//#define DEBUG_KHEAP

#ifndef DEBUG_KHEAP
#define dprintf
#define DBGPRINTF(s)
#endif

// end is defined in the linker script.
extern uint32 end;
uint32 placement_address = (uint32)&end;
extern page_directory_t *kernel_directory;
heap_t *kheap = NULL;

//Physical memory used
volatile uint32 phys_mem_usage;

uint32 kmalloc_int(uint32 sz, int align, uint32 *phys)
{
	static mutex_t mutex;
	mutex_enter(&mutex);
	if (kheap != 0)
	{
		void *addr = alloc(sz, (uint8)align, kheap);
		if (phys != 0)
		{
			page_t *page = get_page((uint32)addr, 0, kernel_directory);
			*phys = page->frame*PAGE_SIZE + ((uint32)addr&0xFFF);
		}
		dprintf ("[kmalloc_int] Alloc %d byte...\n",sz);
		phys_mem_usage += sz;
		mutex_leave(&mutex);
		return (uint32)addr;
	}
	else
	{
		// This will eventually call malloc() on the kernel heap.
		// For now, though, we just assign memory at placement_address
		// and increment it by sz. Even when we've coded our kernel
		// heap, this will be useful for use before the heap is initialised.
		if (align == 1 && (placement_address & 0xFFFFF000) )
		{
			// Align the placement address;
			placement_address &= 0xFFFFF000;
			placement_address += PAGE_SIZE;
		}
		if (phys)
		{
			*phys = placement_address;
		}
		uint32 tmp = placement_address;
		placement_address += sz;
		phys_mem_usage += sz;
		dprintf ("[kmalloc_int] Alloc %d byte...\n",sz);
		mutex_leave(&mutex);
		return tmp;
	}
}

int kfree(void *p)
{
	if(kheap != 0)
	{
		free(p, kheap);
		p = NULL;
		return 0;
	}
	else 
	{
		return 1;
	}
}

uint32 kmalloc_a(uint32 sz)
{
	return kmalloc_int(sz, 1, 0);
}

uint32 kmalloc_p(uint32 sz, uint32 *phys)
{
	return kmalloc_int(sz, 0, phys);
}

uint32 kmalloc_ap(uint32 sz, uint32 *phys)
{
	return kmalloc_int(sz, 1, phys);
}

uint32 kmalloc(uint32 sz)
{
	return kmalloc_int(sz, 0, 0);
}

// В ряде случаев нам нужно будет изменить размер нашей кучи. Если у нас не осталось свободного 
// места, нам потребуется выделить еще памяти. Если мы возвращаем память, нам, возможно,
// потребуется куча меньшего размера.
void expand(uint32 new_size, heap_t *heap)
{
	// Проверяем правильность значений
	ASSERT(new_size > heap->end_address - heap->start_address);
	
	DBGPRINTF('e'); 

	// Получаем следющую ближайшую границу страницы
	if (new_size&0xFFFFF000 != 0)
	{
		new_size &= 0xFFFFF000;
		new_size += PAGE_SIZE;
	}
	
	// Убеждаемся, что мы не перехетрили сами себя
	ASSERT(heap->start_address+new_size <= heap->max_address);
	
	// Это значение должно всегда быть на границе страницы.
	uint32 old_size = heap->end_address-heap->start_address;
	
	uint32 i = old_size;
	while (i < new_size)
	{
		alloc_frame( get_page(heap->start_address+i, 1, kernel_directory),
					(heap->supervisor)?1:0, (heap->readonly)?0:1);
		i += PAGE_SIZE /* page size */;
	}
	heap->end_address = heap->start_address+new_size;
}

uint32 contract(uint32 new_size, heap_t *heap)
{
	DBGPRINTF('c');
	// Проверяем правильность значений
	ASSERT(new_size < heap->end_address-heap->start_address);

	// Получаем следющую ближайшую границу страницы
	if (new_size&PAGE_SIZE)
	{
		new_size &= PAGE_SIZE;
		new_size += PAGE_SIZE;
	}
	
	// Don't contract too far!
	if (new_size < HEAP_MIN_SIZE)
		new_size = HEAP_MIN_SIZE;

	uint32 old_size = heap->end_address-heap->start_address;
	uint32 i = old_size - PAGE_SIZE;
	while (new_size < i)
	{
		free_frame(get_page(heap->start_address+i, 0, kernel_directory));
		i -= PAGE_SIZE;
	}
	
	heap->end_address = heap->start_address + new_size;
	return new_size;
}

// Поиск наименьшего фрагмента, который будет соотвестсвовать определенному количеству байт
static sint32 find_smallest_hole(uint32 size, uint8 page_align, heap_t *heap)
{
	static mutex_t mutex;
	mutex_enter(&mutex);
	// Находим наименьший свободный фрагмент, который подходит
	uint32 iterator = 0;
	while (iterator < heap->index.size)
	{
		header_t *header = (header_t *)lookup_ordered_array(iterator, &heap->index);
		// Если пользователь запросит память, которая выровнена по границе
		if (page_align > 0)
		{
			// Выравниваем по границе начало заголовка
			uint32 location = (uint32)header;
			sint32 offset = 0;
			if ((location+sizeof(header_t)) & 0xFFFFF000 != 0)
				offset = PAGE_SIZE /* page size */  - (location+sizeof(header_t))%PAGE_SIZE;
			sint32 hole_size = (sint32)header->size - offset;
			// Теперь подходит?
			if (hole_size >= (sint32)size)
				break;
		}
		else if (header->size >= size)
			break;
		iterator++;
	}
	mutex_leave(&mutex);
	if (iterator == heap->index.size)
		return -1;					// Мы дошли до конца и ничего не нашли
	else
		return iterator;
}

static sint8 header_t_less_than(void*a, void *b)
{
	return (((header_t*)a)->size < ((header_t*)b)->size)?1:0;
}

heap_t *create_heap(uint32 start, uint32 end_addr, uint32 max, uint8 supervisor, uint8 readonly)
{
	static mutex_t mutex;
	mutex_enter(&mutex);
	heap_t *heap = (heap_t*)kmalloc(sizeof(heap_t));
	
	// Мы предполагаем, что start_address и end_address выровнены по границе страниц
	ASSERT(start%PAGE_SIZE == 0);
	ASSERT(end_addr%PAGE_SIZE == 0);
    
	// Инициализируем список индексов
	heap->index = place_ordered_array( (void*)start, HEAP_INDEX_SIZE, &header_t_less_than);
    
	// Сдвигаем начальный адрес вперед, куда мы сможем начать помещать данные
	start += sizeof(type_t)*HEAP_INDEX_SIZE;

	// Обеспечьте, что бы начальный адрес был выровнен по границе страниц
	if (start & 0xFFFFF000 != 0)
	{
		start &= 0xFFFFF000;
		start += PAGE_SIZE;
	}

	// Запишите начатьный, конечный и максимальный адреса в структуру памяти типа куча
	heap->start_address = start;
	heap->end_address = end_addr;
	heap->max_address = max;
	heap->supervisor = supervisor;
	heap->readonly = readonly;
	
	// Мы начинаем с одного большого фрагьента свободной памяти, указанной в списке индексов
	header_t *hole = (header_t *)start;
	hole->size = end_addr-start;
	hole->magic = HEAP_MAGIC;
	hole->is_hole = 1;
	insert_ordered_array((void*)hole, &heap->index);     
	mutex_leave(&mutex);
	return heap;
}

// Function for allocate memory
void *alloc(uint32 size, uint8 page_align, heap_t *heap)
{
	static mutex_t mutex;
	mutex_enter(&mutex);
	// Проверяем, что мы также учитываем размер заголовка/завершающей записи блока
	uint32 new_size = size + sizeof(header_t) + sizeof(footer_t);
	// Находим наименьший фрагмент свободной памяти, который соответствует нашему запросу.
	sint32 iterator = find_smallest_hole(new_size, page_align, heap);
	
	if (iterator == -1)				// Если не нашелся подходящий фрагмент
	{
		// Сохраняем некоторые ранее полученные данные
		uint32 old_length = heap->end_address - heap->start_address;
		uint32 old_end_address = heap->end_address;
		
		// Нам нужно выделить больше памяти
		expand(old_length+new_size, heap);
		uint32 new_length = heap->end_address-heap->start_address;
		
		// Нахожу самый последний заголовок (по местоположению)
		iterator = 0;
		// В переменных хранятся индекс и значение самого последнего заголовка, найденного к
		// текщему моменту
		uint32 idx = -1; uint32 value = 0x0;
		while (iterator < heap->index.size)
		{
			uint32 tmp = (uint32)lookup_ordered_array(iterator, &heap->index);
			if (tmp > value)
			{
				value = tmp;
				idx = iterator;
			}
			iterator++;
		}
		
		// Если мы не нашли вообще никаких заголовков - нужно 1 заголовок добавить
		if (idx == -1)
		{
			header_t *header = (header_t *)old_end_address;
			header->magic = HEAP_MAGIC;
			header->size = new_length - old_length;
			header->is_hole = 1;
			footer_t *footer = (footer_t *) (old_end_address + header->size - sizeof(footer_t));
			footer->magic = HEAP_MAGIC;
			footer->header = header;
			insert_ordered_array((void*)header, &heap->index);
		}
		else
		{
			// Последний заголовок нужно настроить
			header_t *header = lookup_ordered_array(idx, &heap->index);
			header->size += new_length - old_length;
			// Переписываю завершающую запись
			footer_t *footer = (footer_t *) ((uint32)header + header->size - sizeof(footer_t));
			footer->header = header;
			footer->magic = HEAP_MAGIC;
		}
		mutex_leave(&mutex);
		// Теперь достаточно места, обращаюсь снова к этой функции
		return alloc(size, page_align, heap);
	}
	
	header_t *orig_hole_header = (header_t *)lookup_ordered_array(iterator, &heap->index);
	uint32 orig_hole_pos = (uint32)orig_hole_header;
	uint32 orig_hole_size = orig_hole_header->size;
	// Сюда мы попадаем, если должны разделить фрагмент, который мы нашли, на две части.
	// Является ли размер запрашиваемого фрагмента меньше накладных расходов нв 
	// добавление нового фрагмента?
	if (orig_hole_size-new_size < sizeof(header_t)+sizeof(footer_t))
	{
		// Тогда просто увеличиваем запрашиваемый размер до размера найденного фрагмента.
		size += orig_hole_size-new_size;
		new_size = orig_hole_size;
	}
	
	// Если нам нужно выравнивать данные по границе страницы, то выполняем это сейчас
	// и делаем так, чтобы новый фрагмент находился перед нашим блоком памяти
	if (page_align && orig_hole_pos&0xFFFFF000)
	{
		uint32 new_location   = orig_hole_pos + PAGE_SIZE /* page size */ - (orig_hole_pos&0xFFF) - sizeof(header_t);
		header_t *hole_header = (header_t *)orig_hole_pos;
		hole_header->size     = PAGE_SIZE /* page size */ - (orig_hole_pos&0xFFF) - sizeof(header_t);
		hole_header->magic    = HEAP_MAGIC;
		hole_header->is_hole  = 1;
		footer_t *hole_footer = (footer_t *) ( (uint32)new_location - sizeof(footer_t) );
		hole_footer->magic    = HEAP_MAGIC;
		hole_footer->header   = hole_header;
		orig_hole_pos         = new_location;
		orig_hole_size        = orig_hole_size - hole_header->size;
	}
	else
	{
		// В противном случае этот фрагмент нам больше не нужен, удаляю его из списка индексов.
		remove_ordered_array(iterator, &heap->index);
	}

	// Перезаписываем исходный заголовок...
	header_t *block_header  = (header_t *)orig_hole_pos;
	block_header->magic     = HEAP_MAGIC;
	block_header->is_hole   = 0;
	block_header->size      = new_size;
	// ... И завершающую запись
	footer_t *block_footer  = (footer_t *) (orig_hole_pos + sizeof(header_t) + size);
	block_footer->magic     = HEAP_MAGIC;
	block_footer->header    = block_header;

	// Нам после выделенного блока может потребоваться записать новый фрагмент
	// Мы делаем это только если размер нового фрагмента положительный ...
	if (orig_hole_size - new_size > 0)
	{
		header_t *hole_header = (header_t *) (orig_hole_pos + sizeof(header_t) + size + sizeof(footer_t));
		hole_header->magic    = HEAP_MAGIC;
		hole_header->is_hole  = 1;
		hole_header->size     = orig_hole_size - new_size;
		footer_t *hole_footer = (footer_t *) ( (uint32)hole_header + orig_hole_size - new_size - sizeof(footer_t) );
		if ((uint32)hole_footer < heap->end_address)
		{
			hole_footer->magic = HEAP_MAGIC;
			hole_footer->header = hole_header;
		}
		// Помещаем новый фрагмент в индексный список
		insert_ordered_array((void*)hole_header, &heap->index);
	}
	mutex_leave(&mutex);
	// Все!
	return (void *) ( (uint32)block_header+sizeof(header_t) );
}

void free(void *p, heap_t *heap)
{
	static mutex_t mutex;
	// Выход в случае нулевых указателей
	if (p == NULL)
		return;
	mutex_enter(&mutex);
	
	// Берем заголовок и завершающую запись, связанные с этим указателем
	header_t *header = (header_t*) ( (uint32)p - sizeof(header_t) );
	footer_t *footer = (footer_t*) ( (uint32)header + header->size - sizeof(footer_t) );
	
	phys_mem_usage -= header->size;
	dprintf ("[free] free %d byte...\n", header->size);

	// Проверяем правильность значений
	ASSERT(header->magic == HEAP_MAGIC);
	ASSERT(footer->magic == HEAP_MAGIC);

	// Создаем фрагмент свободной памяти
	header->is_hole = 1;

	// Хотим ли мы добавить этот заголовок в список индексов "фрагменты свободной памяти" index?
	char do_add = 1;

	// Унификация слева
	// Если непосредственно слева от фрагмента находится заверша.щая запись другого фрагмента ...
	footer_t *test_footer = (footer_t*) ((uint32)header - sizeof(footer_t));
	if (test_footer->magic == HEAP_MAGIC &&
			test_footer->header->is_hole == 1)
	{
		uint32 cache_size = header->size;				// Кешируем наш текущий размер
		header = test_footer->header;					// Перезаписываем вместо нашего новый заголовок
		footer->header = header;					// Заменяем нашу завершающую запись новой, на которую указывает заголовок
		header->size += cache_size;					// Изменяем размер
		do_add = 0;							// Поскольку этот заголовок уже в индексном списке снова не добавляю
	}

	// Унификация справа
	// Если непосредственно справа от фрагмента находится заголовок другого фрагмента ...
	header_t *test_header = (header_t*) ((uint32)footer + sizeof(footer_t));
	if (test_header->magic == HEAP_MAGIC &&
			test_header->is_hole)
	{
		header->size += test_header->size;				// Увеличиваю размер
		test_footer = (footer_t*) ((uint32)test_header+	// Переписываю завершающую запись на которую указывает наш заголовок
						test_header->size - sizeof(footer_t));
		footer = test_footer;
		
		// Находим и удаляем этот заголовок из списка индексов
		uint32 iterator = 0;
		while ((iterator < heap->index.size) &&
					(lookup_ordered_array(iterator, &heap->index) != (void*)test_header))
			iterator++;

		// Убеждаюсь, что я нашел правильный объект
		dprintf ("[free] ASSERT iterator = %d heap->index.size = %d...\n", iterator,heap->index.size);
		//ASSERT(iterator < heap->index.size);
		// Проверка на то что мы что-то нашли... Удаляю его
		if (iterator < heap->index.size)
			remove_ordered_array(iterator, &heap->index);
	}

	// Если положение завершающей записи совпадает с конечным адресом кучи - можно уменьшить кучу
	if ( (uint32)footer+sizeof(footer_t) == heap->end_address)
	{
		uint32 old_length = heap->end_address-heap->start_address;
		uint32 new_length = contract((uint32)header - heap->start_address, heap);
		// Проверяем насколько большой будет куча после изменения размера
		if (header->size - (old_length-new_length) > 0)
		{
			// Все еще существуем, теперь изменяю размер
			header->size -= old_length-new_length;
			footer = (footer_t*) ((uint32)header + header->size - sizeof(footer_t));
			footer->magic = HEAP_MAGIC;
			footer->header = header;
		}
		else
		{
			// Уже нет, удаляю из списка индексов
			uint32 iterator = 0;
			while ((iterator < heap->index.size) &&
					(lookup_ordered_array(iterator, &heap->index) != (void*)test_header))
				iterator++;
			// Если не нашел - удалять нечего
			if (iterator < heap->index.size)
				remove_ordered_array(iterator, &heap->index);
		}
	}
	
	if (do_add == 1)
		insert_ordered_array((void*)header, &heap->index);
	mutex_leave(&mutex);
}

uint32 get_memory_usage()
{
    return phys_mem_usage+(placement_address-0x1000); 
}